package com.starhub.application.rag.biz;

import com.starhub.application.rag.service.MilvusConfigService;
import com.starhub.common.bean.model.MilvusConfig;
import com.starhub.common.sys.user.LoginUser;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import jakarta.annotation.PreDestroy;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @Description:(动态加载milvus配置)
 * @version: V1.0
 * @author: liuhf
 */
@Slf4j
@Component
public class MilvusClientFactory {

    @Autowired
    private MilvusConfigService milvusConfigService;

    private final Map<String, ReentrantLock> lockMap = new ConcurrentHashMap<>();
    
    // 客户端连接池，用于缓存和复用连接
    private final Map<String, MilvusServiceClient> clientPool = new ConcurrentHashMap<>();
    
    // 客户端使用状态跟踪
    private final Map<String, Long> clientLastUsedTime = new ConcurrentHashMap<>();

    private static final int MAX_RETRY_TIMES = 3;
    private static final long RETRY_INTERVAL_MS = 1000; // 重试间隔1秒
    private static final long CLIENT_IDLE_TIMEOUT_MS = 5 * 60 * 1000; // 客户端空闲超时时间：5分钟

    /**
     * 创建Milvus客户端
     * @param config Milvus配置
     * @return Milvus客户端实例
     */
    private MilvusServiceClient createClient(MilvusConfig config) {
        int retryCount = 0;
        Exception lastException = null;

        while (retryCount < MAX_RETRY_TIMES) {
            try {
                ConnectParam connectParam = ConnectParam.newBuilder()
                        .withHost(config.getHost())
                        .withPort(config.getPort())
                        .withAuthorization(config.getUsername(), config.getPassword())
                        .withSecure(false) // 根据配置决定是否使用SSL
                        .withDatabaseName(config.getDatabaseName())
                        .withConnectTimeout(10, TimeUnit.SECONDS)  // 增加连接超时时间
                        .withKeepAliveTime(30, TimeUnit.SECONDS)
                        .withKeepAliveTimeout(10, TimeUnit.SECONDS)
                        .withIdleTimeout(60, TimeUnit.SECONDS)
                        .build();

                return new MilvusServiceClient(connectParam);
            } catch (Exception e) {
                lastException = e;
                retryCount++;
                if (retryCount < MAX_RETRY_TIMES) {
                    log.warn("创建Milvus客户端失败，正在进行第{}次重试: {}", retryCount, e.getMessage());
                    try {
                        Thread.sleep(RETRY_INTERVAL_MS);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new RuntimeException("重试过程被中断", ie);
                    }
                }
            }
        }

        log.error("创建Milvus客户端失败，已重试{}次: {}", MAX_RETRY_TIMES, lastException.getMessage());
        throw new RuntimeException("创建Milvus客户端失败，已重试" + MAX_RETRY_TIMES + "次: " + lastException.getMessage());
    }

    /**
     * 获取Milvus客户端（带连接池管理）
     * @param mark 标记
     * @return Milvus客户端实例
     */
    public MilvusServiceClient getClient(String mark) {
        ReentrantLock lock = lockMap.computeIfAbsent(mark, k -> new ReentrantLock());
        lock.lock();
        try {
            MilvusConfig config = milvusConfigService.getConfigByMark(mark);
            if (config == null) {
                throw new RuntimeException("未找到mark为" + mark + "的Milvus配置");
            }
            
            // 检查连接池中是否有可用的客户端
            MilvusServiceClient client = clientPool.get(mark);
            if (client != null) {
                // 检查客户端是否仍然有效
                try {
                    // 简单的健康检查 - 使用一个简单的操作来验证连接
                    client.listCollections(null);
                    clientLastUsedTime.put(mark, System.currentTimeMillis());
                    return client;
                } catch (Exception e) {
                    log.warn("连接池中的客户端已失效，将创建新客户端: {}", e.getMessage());
                    closeClient(mark);
                }
            }
            
            // 创建新客户端
            client = createClient(config);
            clientPool.put(mark, client);
            clientLastUsedTime.put(mark, System.currentTimeMillis());
            return client;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 安全关闭客户端
     * @param mark 标记
     */
    public void closeClient(String mark) {
        MilvusServiceClient client = clientPool.remove(mark);
        if (client != null) {
            try {
                client.close();
                log.info("成功关闭Milvus客户端: {}", mark);
            } catch (Exception e) {
                log.warn("关闭Milvus客户端时发生异常: {}", e.getMessage());
            }
        }
        clientLastUsedTime.remove(mark);
    }

    /**
     * 关闭所有客户端连接
     */
    @PreDestroy
    public void closeAllClients() {
        log.info("开始关闭所有Milvus客户端连接...");
        for (String mark : clientPool.keySet()) {
            closeClient(mark);
        }
        log.info("所有Milvus客户端连接已关闭");
    }

    /**
     * 清理空闲的客户端连接
     */
    public void cleanupIdleClients() {
        long currentTime = System.currentTimeMillis();
        for (Map.Entry<String, Long> entry : clientLastUsedTime.entrySet()) {
            String mark = entry.getKey();
            Long lastUsedTime = entry.getValue();
            
            if (currentTime - lastUsedTime > CLIENT_IDLE_TIMEOUT_MS) {
                log.info("清理空闲的Milvus客户端: {}", mark);
                closeClient(mark);
            }
        }
    }

    /**
     * 添加配置
     * @param config Milvus配置
     */
    public void addConfig(MilvusConfig config) {
        milvusConfigService.add(null, config);
    }

    /**
     * 删除配置
     * @param mark 标记
     */
    public void deleteConfig(String mark) {
        MilvusConfig config = milvusConfigService.getConfigByMark(mark);
        if (config != null) {
            milvusConfigService.deleteByType(null, config.getId().toString(), 1);
            // 删除配置时同时关闭对应的客户端
            closeClient(mark);
        }
    }

    /**
     * 更新配置
     * @param config Milvus配置
     */
    public void updateConfig(MilvusConfig config) {
        milvusConfigService.update(null, config);
        // 更新配置时关闭旧的客户端，下次获取时会创建新的
        closeClient(config.getMark());
    }

    /**
     * 校验某个连接是否有效
     * @param user 当前用户
     * @param mark 标记
     * @return 是否有效
     */
    public boolean validateConnection(LoginUser user, String mark) {
        ReentrantLock lock = lockMap.computeIfAbsent(mark, k -> new ReentrantLock());
        lock.lock();
        try {
            MilvusConfig config = milvusConfigService.getConfigByMark(mark);
            if (config == null) {
                return false;
            }
            try {
                MilvusServiceClient client = createClient(config);
                client.close();
                return true;
            } catch (Exception e) {
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 获取所有配置
     * @return 配置列表
     */
    public List<MilvusConfig> getAllConfigs() {
        return milvusConfigService.getAllConfigs();
    }
}